"""高级技术指标计算函数."""

from __future__ import annotations

import numpy as np
import pandas as pd
from scipy import stats

from .utils import _ema, _sma


def _vwap(high: pd.Series, low: pd.Series, close: pd.Series, volume: pd.Series) -> pd.Series:
    """计算 VWAP（成交量加权平均价）.
    
    VWAP是机构投资者常用的价格基准，通过成交量加权计算平均价格，
    能够反映真实的交易成本。
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    volume:
        成交量序列
        
    Returns
    -------
    pd.Series
        VWAP序列
    """
    # 计算典型价格（Typical Price）：(最高价 + 最低价 + 收盘价) / 3
    typical_price = (high + low + close) / 3
    
    # 计算成交量加权价格：典型价格 × 成交量
    typical_price_volume = typical_price * volume
    
    # 计算累计成交量加权价格和累计成交量
    cumulative_tpv = typical_price_volume.cumsum()
    cumulative_volume = volume.cumsum()
    
    # VWAP = 累计成交量加权价格 / 累计成交量
    vwap = cumulative_tpv / cumulative_volume.replace(0, np.nan)
    
    return vwap.replace([np.inf, -np.inf], np.nan)


def _parkinson_volatility(high: pd.Series, low: pd.Series, window: int = 20) -> pd.Series:
    """计算 Parkinson 波动率.
    
    Parkinson波动率基于日内高低价，比基于收盘价的传统波动率估计更准确，
    因为它利用了更多的价格信息。
    
    公式：σ_p = sqrt((1/(4*ln(2))) * (ln(high/low))^2)
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    window:
        滚动窗口大小（默认20日）
        
    Returns
    -------
    pd.Series
        Parkinson波动率序列
    """
    # 计算 ln(high/low)
    log_hl = np.log(high / low.replace(0, np.nan))
    
    # 计算 (ln(high/low))^2
    log_hl_squared = log_hl ** 2
    
    # Parkinson波动率公式：sqrt((1/(4*ln(2))) * mean((ln(high/low))^2))
    # 常数 1/(4*ln(2)) ≈ 0.361
    parkinson_constant = 1.0 / (4.0 * np.log(2.0))
    
    # 使用滚动窗口计算均值
    parkinson_vol = np.sqrt(parkinson_constant * log_hl_squared.rolling(window, min_periods=window).mean())
    
    return parkinson_vol.replace([np.inf, -np.inf], np.nan)


def _price_percentile_rank(close: pd.Series, window: int) -> pd.Series:
    """计算价格百分位排名.
    
    计算当前价格在过去N日价格区间中的百分位排名，反映价格相对位置。
    百分位排名范围0-100，值越大表示价格越接近历史高位。
    
    Parameters
    ----------
    close:
        收盘价序列
    window:
        滚动窗口大小（如20、60、120日）
        
    Returns
    -------
    pd.Series
        价格百分位排名序列（0-100）
    """
    # 使用滚动窗口计算百分位排名
    # 对于每个时间点，计算当前价格在该窗口内所有价格中的百分位排名
    percentile_rank = close.rolling(window, min_periods=window).apply(
        lambda x: stats.percentileofscore(x.iloc[:-1], x.iloc[-1], kind='rank') if len(x) > 1 else np.nan,
        raw=False
    )
    return percentile_rank


def _roc(close: pd.Series, window: int) -> pd.Series:
    """计算 ROC（变动率指标，Rate of Change）.
    
    ROC反映价格变化率，与RSI互补，用于衡量价格动量。
    公式：ROC = (当前价格 - N日前价格) / N日前价格 × 100
    
    Parameters
    ----------
    close:
        收盘价序列
    window:
        窗口大小（如10、20日）
        
    Returns
    -------
    pd.Series
        ROC序列（百分比）
    """
    # ROC = (当前价格 - N日前价格) / N日前价格 × 100
    # 等价于 pct_change * 100，但明确标注为ROC
    roc = close.pct_change(window) * 100
    return roc.replace([np.inf, -np.inf], np.nan)


def _volatility_cone(volatility_series: pd.Series, lookback_window: int = 252) -> pd.Series:
    """计算波动率锥（Volatility Cone）.
    
    计算当前波动率在历史波动率分布中的百分位排名，用于判断当前波动率水平是否异常。
    波动率锥范围0-100，值越大表示当前波动率越接近历史高位。
    
    Parameters
    ----------
    volatility_series:
        波动率序列
    lookback_window:
        历史窗口大小（默认252日，约1年）
        
    Returns
    -------
    pd.Series
        波动率锥序列（0-100）
    """
    # 使用滚动窗口计算波动率在历史分布中的百分位排名
    volatility_cone = volatility_series.rolling(lookback_window, min_periods=lookback_window).apply(
        lambda x: stats.percentileofscore(x.iloc[:-1], x.iloc[-1], kind='rank') if len(x) > 1 else np.nan,
        raw=False
    )
    return volatility_cone


def _trix(close: pd.Series, window: int = 14) -> pd.Series:
    """计算 TRIX（三重指数平滑移动平均）.
    
    TRIX是长期趋势指标，通过三重EMA平滑过滤短期波动，适合捕捉长期趋势变化。
    计算步骤：
    1. EMA1 = EMA(close, window)
    2. EMA2 = EMA(EMA1, window)
    3. EMA3 = EMA(EMA2, window)
    4. TRIX = (EMA3 - EMA3.shift(1)) / EMA3.shift(1) × 100
    
    Parameters
    ----------
    close:
        收盘价序列
    window:
        窗口大小（默认14日）
        
    Returns
    -------
    pd.Series
        TRIX序列（百分比）
    """
    # 第一重EMA
    ema1 = _ema(close, window)
    # 第二重EMA
    ema2 = _ema(ema1, window)
    # 第三重EMA
    ema3 = _ema(ema2, window)
    # TRIX = (EMA3 - EMA3.shift(1)) / EMA3.shift(1) × 100
    trix = ema3.pct_change(1) * 100
    return trix.replace([np.inf, -np.inf], np.nan)


def _trend_persistence(close: pd.Series, ma: pd.Series) -> pd.Series:
    """计算趋势持续性.
    
    计算价格在移动平均线上方/下方的连续天数，反映趋势的持续性。
    正数表示价格在均线上方，负数表示在下方，绝对值表示连续天数。
    
    Parameters
    ----------
    close:
        收盘价序列
    ma:
        移动平均线序列（如SMA20、EMA20）
        
    Returns
    -------
    pd.Series
        趋势持续性序列（正数表示上方，负数表示下方）
    """
    # 判断价格是否在均线上方（True表示上方，False表示下方）
    above_ma = close > ma
    
    # 计算连续天数：使用groupby识别连续相同值的组
    # 创建分组标识：当above_ma的值改变时，创建新的组
    groups = (above_ma != above_ma.shift(1)).cumsum()
    
    # 计算每个组的累计计数
    persistence = pd.Series(index=close.index, dtype=float)
    
    for group_id in groups.unique():
        mask = groups == group_id
        group_size = mask.sum()
        is_above = above_ma[mask].iloc[0] if group_size > 0 else False
        
        # 为组内每个位置分配连续天数（从1开始，保持符号）
        for i, idx in enumerate(close.index[mask], start=1):
            persistence.loc[idx] = i if is_above else -i
    
    return persistence


def _garman_klass_volatility(open: pd.Series, high: pd.Series, low: pd.Series, close: pd.Series, window: int = 20) -> pd.Series:
    """计算 Garman-Klass 波动率.
    
    Garman-Klass波动率结合OHLC四个价格的波动率估计，比Parkinson波动率更准确，
    能够更充分利用价格信息。
    
    公式：GK = sqrt(0.5 * (ln(high/low))^2 - (2*ln(2)-1) * (ln(close/open))^2)
    
    Parameters
    ----------
    open:
        开盘价序列
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    window:
        滚动窗口大小（默认20日）
        
    Returns
    -------
    pd.Series
        Garman-Klass波动率序列
    """
    # 计算 ln(high/low)
    log_hl = np.log(high / low.replace(0, np.nan))
    log_hl_squared = log_hl ** 2
    
    # 计算 ln(close/open)
    log_co = np.log(close / open.replace(0, np.nan))
    log_co_squared = log_co ** 2
    
    # Garman-Klass波动率公式
    # 常数 (2*ln(2)-1) ≈ 0.386
    constant = 2 * np.log(2) - 1
    
    # 计算每日的Garman-Klass波动率
    gk_daily = np.sqrt(0.5 * log_hl_squared - constant * log_co_squared)
    
    # 使用滚动窗口计算均值
    gk_vol = gk_daily.rolling(window, min_periods=window).mean()
    
    return gk_vol.replace([np.inf, -np.inf], np.nan)

